import cv2
import numpy as np
from typing import Optional, Tuple
import time
from threading import Thread, Event
from queue import Queue
import traceback


class StreamProcessor:
    def __init__(self, rtsp_url: str, target_fps: int = 10):
        self.rtsp_url = rtsp_url                # RTSP URL
        self.target_fps = target_fps            # 目标帧率
        self.frame_interval = 1.0 / target_fps  # 帧间隔

        self.cap = None                         # opencv视频捕获对象
        self.frame_queue = Queue(maxsize=30)    # 帧队列，用于存放最新帧队列，最多30帧
        self.is_running = False                 # 流处理器是否运行
        self.capture_thread = None              # 捕获帧线程对象
        self.last_frame_time = 0                # 上次捕获帧的时间戳

    def start(self) -> bool:
        """启动流处理"""
        try:
            self.cap = cv2.VideoCapture(self.rtsp_url)          # 创建视频捕获对象rtsp流，使用opencv的VideoCapture
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)      # 设置缓冲区大小为1，减少缓冲延迟

            if not self.cap.isOpened():                         # rtsp视频流是否打开
                print(f"无法打开RTSP流: {self.rtsp_url}")
                return False

            self.is_running = True                                      # 设置流处理器为运行状态
            self.capture_thread = Thread(target=self._capture_frames)   # 创建捕获帧线程，目标为捕获帧线程函数（）
            self.capture_thread.daemon = True                           # 设置捕获帧线程为守护线程
            self.capture_thread.start()                                 # 启动捕获帧线程

            print(f"成功连接到RTSP流: {self.rtsp_url}")
            return True

        except Exception as e:
            print(f"启动流处理失败: {e}")
            return False

    def _capture_frames(self):
        """捕获帧的线程函数"""
        while self.is_running:
            try:
                ret, frame = self.cap.read()    # 读取帧，ret为是否成功，frame为帧
                if ret:
                    current_time = time.time()  # 获取当前时间戳

                    # 控制帧率：只有距离上次取帧时间超过帧间隔才处理
                    if current_time - self.last_frame_time >= self.frame_interval:
                        # 清空旧帧，只保留最新帧，降低延迟
                        while not self.frame_queue.empty():
                            try:
                                self.frame_queue.get_nowait()   # 从队列中获取帧，不等待
                            except:
                                pass

                        self.frame_queue.put(frame)             # 将帧添加到队列，队列满时，旧帧会被自动丢弃
                        self.last_frame_time = current_time     # 更新上次取帧时间
                else:
                    print("读取帧失败，尝试重连...")
                    self._reconnect()

            except Exception as e:
                print(f"捕获帧错误: {e}")
                time.sleep(0.1)

    def _reconnect(self):
        """重连RTSP流"""
        try:
            if self.cap:
                self.cap.release()  # 释放视频捕获对象
            time.sleep(1)           # 等待1秒
            self.cap = cv2.VideoCapture(self.rtsp_url)      # 重新创建视频捕获对象
            self.cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)  # 设置缓冲区大小为1
        except Exception as e:
            print(f"重连失败: {e}")

    def get_frame(self) -> Optional[np.ndarray]:
        """获取最新帧"""
        try:
            if not self.frame_queue.empty():                # 如果队列不为空
                return self.frame_queue.get(timeout=0.1)    # 从队列中获取帧，timeout=0.1秒
            return None                                     # 如果队列为空，返回None
        except:
            return None

    def stop(self):
        """停止流处理"""
        self.is_running = False     # 设置流处理器为停止状态
        if self.capture_thread:     # 如果捕获帧线程存在
            self.capture_thread.join(timeout=2)  # 等待捕获帧线程结束，最多等待2秒
        if self.cap:                # 如果视频捕获对象存在
            self.cap.release()      # 释放视频捕获对象
        print("流处理已停止")
